# -*- coding: utf-8 -*-

import os
from os.path import dirname

import h5py
import numpy as np
import time
from pprint import pprint
from milvus import Milvus, IndexType, MetricType

THRESHOLD = float(os.environ.get('THRESHOLD', '0.85'))  # 检索阈值

class MilvusRetrieval(object):
    def __init__(self, index_name, index_dir,
        host=os.environ.get("MILVUS_HOST", "127.0.0.1"),
        port=os.environ.get("MILVUS_PORT", 19530)):

        print(f"[INFO] Connecting to Milvus at {host}:{port} ...")
        t0 = time.time()
        self.client = Milvus(host, port)
        print(f"[INFO] ✅ Connected to Milvus ({time.time() - t0:.3f}s)\n")

        self.index_name = index_name
        self.load(index_dir)

    def load(self, index_dir):
        print(f"[INFO] Loading index data from: {index_dir}")
        t0 = time.time()

        # 1. 读取索引文件
        h5f = h5py.File(index_dir, 'r')
        self.retrieval_db = h5f['dataset_1'][:]
        self.retrieval_name = h5f['dataset_2'][:]
        h5f.close()
        print(f"[INFO] 读取索引文件完成，用时: {time.time() - t0:.3f}s")
        print(f"[INFO] 向量数量: {len(self.retrieval_db)}")

        # 预览部分数据

        print(f"[INFO] 读取索引文件完成，用时: {time.time() - t0:.3f}s")
        print(f"[INFO] 向量数量: {len(self.retrieval_db)}")
        print(f"[INFO] 向量维度: {self.retrieval_db.shape[1] if len(self.retrieval_db.shape) > 1 else '未知'}")

        # 2. 删除旧集合（若存在）
        if self.index_name in self.client.list_collections()[1]:
            print(f"[INFO] 发现已有集合 '{self.index_name}'，正在删除...")
            t_del = time.time()
            self.client.drop_collection(collection_name=self.index_name)
            print(f"[INFO] 删除旧集合完成，用时: {time.time() - t_del:.3f}s")
            time.sleep(3)#删除集合属于异步操作，等待删除完毕再执行后续内容

        # 3. 创建新集合
        print("[INFO] 创建新集合...")
        t_create = time.time()
        self.client.create_collection({
            'collection_name': self.index_name,
            'dimension': 512,
            'index_file_size': 1024,
            'metric_type': MetricType.IP
        })
        print(f"[INFO] 创建集合完成，用时: {time.time() - t_create:.3f}s")

        # 4. 插入向量
        print("[INFO] 正在插入向量到 Milvus ...")
        t_insert = time.time()
        status, ids = self.client.insert(
            collection_name=self.index_name,
            records=[i.tolist() for i in self.retrieval_db]
        )
        print(f"[INFO] 插入完成，用时: {time.time() - t_insert:.3f}s")

        # 5. 建立 ID 映射关系
        self.id_dict = {}
        for i, val in enumerate(self.retrieval_name):
            if i < len(ids):
                self.id_dict[ids[i]] = str(val)
            else:
                print(f"[WARNING] ID 数量不足: i={i}, len(ids)={len(ids)}")
                break

        # 6. 创建索引
        """目的：加速高维向量相似检索（Approximate Nearest Neighbor Search）
         适用场景：非结构化数据（图像、语音、文本、特征向量）,作用：在百万甚至上亿个向量中快速找到最相似的若干个,实现方式：使用 高维空间的近似搜索算法，如：
            FLAT（暴力全量搜索）, IVF_FLAT, HNSW,ANNOY,R-Tree / PQ 等
        """
        print("[INFO] 创建索引 (IndexType.FLAT) ...")
        t_index = time.time()
        self.client.create_index(self.index_name, IndexType.FLAT, {'nlist': 16384})
        print(f"[INFO] 索引创建完成，用时: {time.time() - t_index:.3f}s")

        print(f"[INFO] ✅ 索引加载完毕，总耗时: {time.time() - t0:.3f}s\n")
        print(f"************* Done milvus indexing, Indexed {len(self.retrieval_db)} documents *************")

    def retrieve(self, query_vector, search_size=3):
        print(f"[INFO] 开始检索 top-{search_size} 向量...")
        t0 = time.time()

        r_list = []
        status, vectors = self.client.search(
            collection_name=self.index_name,
            query_records=[query_vector],
            top_k=search_size,
            params={'nprobe': 16}
        )

        search_time = time.time() - t0
        print(f"[INFO] 检索完成，用时: {search_time:.3f}s")

        # 结果过滤
        for v in vectors[0]:
            score = float(v.distance) * 0.5 + 0.5
            if score > THRESHOLD:
                temp = {
                    "id": v.id,
                    "name": self.id_dict.get(v.id, "unknown"),
                    "score": round(score, 6)
                }
                r_list.append(temp)

        print(f"[INFO] 返回结果数量: {len(r_list)} (阈值: {THRESHOLD})\n")
        return r_list


